package zkp

import (
	"os"

	"github.com/consensys/gnark-crypto/ecc"
	"github.com/consensys/gnark/backend/groth16"
	"github.com/consensys/gnark/frontend"
	"github.com/consensys/gnark/frontend/cs/r1cs"
)

const (
	provingKeyFile   = "auth_proving.key"
	verifyingKeyFile = "auth_verifying.key"
)

// KeyManager 管理证明和验证密钥
type KeyManager struct {
	provingKey   groth16.ProvingKey
	verifyingKey groth16.VerifyingKey
}

// NewKeyManager 创建新的密钥管理器
func NewKeyManager() (*KeyManager, error) {
	km := &KeyManager{}

	// 检查密钥文件是否存在
	if _, err := os.Stat(provingKeyFile); os.IsNotExist(err) {
		if err := km.generateKeys(); err != nil {
			return nil, err
		}
	} else {
		if err := km.loadKeys(); err != nil {
			return nil, err
		}
	}

	return km, nil
}

// generateKeys 生成新的密钥对
func (km *KeyManager) generateKeys() error {
	// 创建空电路
	circuit := &AuthCircuit{}

	// 编译电路
	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, circuit)
	if err != nil {
		return err
	}

	// 生成证明密钥和验证密钥
	pk, vk, err := groth16.Setup(ccs)
	if err != nil {
		return err
	}

	// 保存密钥到文件
	if err := saveKey(pk, provingKeyFile); err != nil {
		return err
	}

	if err := saveKey(vk, verifyingKeyFile); err != nil {
		return err
	}

	km.provingKey = pk
	km.verifyingKey = vk

	return nil
}

// loadKeys 从文件加载密钥
func (km *KeyManager) loadKeys() error {
	pk, err := loadProvingKey(provingKeyFile)
	if err != nil {
		return err
	}

	vk, err := loadVerifyingKey(verifyingKeyFile)
	if err != nil {
		return err
	}

	km.provingKey = pk
	km.verifyingKey = vk

	return nil
}

// GetProvingKey 获取证明密钥
func (km *KeyManager) GetProvingKey() groth16.ProvingKey {
	return km.provingKey
}

// GetVerifyingKey 获取验证密钥
func (km *KeyManager) GetVerifyingKey() groth16.VerifyingKey {
	return km.verifyingKey
}

// 辅助函数：保存密钥到文件
func saveKey(key interface{}, filename string) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	switch k := key.(type) {
	case groth16.ProvingKey:
		_, err = k.WriteTo(file)
	case groth16.VerifyingKey:
		_, err = k.WriteTo(file)
	}

	return err
}

// 辅助函数：从文件加载证明密钥
func loadProvingKey(filename string) (groth16.ProvingKey, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	pk := groth16.NewProvingKey(ecc.BN254)
	_, err = pk.ReadFrom(file)
	return pk, err
}

// 辅助函数：从文件加载验证密钥
func loadVerifyingKey(filename string) (groth16.VerifyingKey, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	vk := groth16.NewVerifyingKey(ecc.BN254)
	_, err = vk.ReadFrom(file)
	return vk, err
}
